#ifndef THREE_N_N_N_D_H
#define THREE_N_N_N_D_H

#include "kernel_operator.h"

namespace ThreeNN {
using namespace AscendC;

constexpr int32_t BYTE_BLOCK = 32;

constexpr float FLOAT_MAX = 4e38;

template <typename T, int32_t BUFFER_NUM>
class ThreeNNND {
public:
    __aicore__ inline ThreeNNND() = default;
    __aicore__ inline void Init(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices,
                                const ThreeNNTilingData* __restrict tilingData);
    __aicore__ inline void ProcessSmallX2();
    __aicore__ inline void ProcessBigX2();

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyInX2(int64_t gmOffset, int64_t dataCount);
    __aicore__ inline void CopySplit(int64_t groupCount);

    __aicore__ inline void ComputePerf(T x1, T y1, T z1, int64_t groupCount);
    __aicore__ inline void CmpBest(T tempData, int32_t tempIndex);
    __aicore__ inline void ComputeNormal(T x1, T y1, T z1, int64_t groupCount, int64_t indexOffset);

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> x2Queue;
    GlobalTensor<T> x1GM, x2GM, distGM;
    GlobalTensor<int32_t> idxGM;
    TBuf<QuePosition::VECCALC> tempValBuf;
    LocalTensor<T> tempValLT;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;
    uint64_t ubMaxGroupCount = 0;
    uint64_t tailGroupCount = 0;
    T best1 = FLOAT_MAX;
    T best2 = FLOAT_MAX;
    T best3 = FLOAT_MAX;
    int32_t bestIndex1 = 0;
    int32_t bestIndex2 = 0;
    int32_t bestIndex3 = 0;

    // tiling params
    uint64_t ubMaxProcCount = 0;
    uint64_t bSize = 0;
    uint64_t nSize = 0;
    uint64_t mSize = 0;
    uint64_t loopTime = 0;
    uint64_t tailCount = 0;
};

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::Init(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices,
                                                      const ThreeNNTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);
    ubMaxProcCount = tilingData->ubMaxProcCount;
    bSize = tilingData->bSize;
    nSize = tilingData->nSize;
    mSize = tilingData->mSize;
    loopTime = tilingData->loopTime;
    tailCount = tilingData->tailCount;

    ubMaxGroupCount = ubMaxProcCount / 3;
    tailGroupCount = tailCount / 3;
    uint64_t x1DataCount = bSize * nSize * 3;
    uint64_t x2DataCount = bSize * mSize * 3;

    x1GM.SetGlobalBuffer((__gm__ T*)xyz1, x1DataCount);
    x2GM.SetGlobalBuffer((__gm__ T*)xyz2, x2DataCount);
    distGM.SetGlobalBuffer((__gm__ T*)dist, x1DataCount);
    idxGM.SetGlobalBuffer((__gm__ int32_t*)indices, x1DataCount);

    uint64_t singleBufferSize = ubMaxProcCount * sizeof(T);
    pipe.InitBuffer(x2Queue, BUFFER_NUM, singleBufferSize);
    pipe.InitBuffer(tempValBuf, singleBufferSize * (3 - BUFFER_NUM));
    tempValLT = tempValBuf.Get<T>();

    // printf(
    //     "==zf==Kernel Tiling Data, ubMaxProcCount:%lu, ubMaxGroupCount:%lu, tailGroupCount:%lu, bSize:%lu, nSize:%lu, "
    //     "mSize:%lu, "
    //     "loopTime:%lu, tailCount:%lu\n",
    //     ubMaxProcCount, ubMaxGroupCount, tailGroupCount, bSize, nSize, mSize, loopTime, tailCount);
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::ProcessSmallX2() {
    T x1, y1, z1;
    int64_t gmOffset = 0, singleAlignCount = CeilAlignA2B(mSize * 3, perBlockCount);
    for (int64_t i = 0; i < bSize; i++) {
        CopyInX2(gmOffset, singleAlignCount);
        CopySplit(mSize);
        for (int64_t j = 0, index = i * nSize * 3; j < nSize; j++, index += 3) {
            x1 = x1GM.GetValue(index);
            y1 = x1GM.GetValue(index + 1);
            z1 = x1GM.GetValue(index + 2);
            ComputePerf(x1, y1, z1, mSize);
            distGM.SetValue(index, best1);
            distGM.SetValue(index + 1, best2);
            distGM.SetValue(index + 2, best3);
            idxGM.SetValue(index, bestIndex1);
            idxGM.SetValue(index + 1, bestIndex2);
            idxGM.SetValue(index + 2, bestIndex3);
        }
        gmOffset += mSize * 3;
    }
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::ProcessBigX2() {
    int64_t gmOffsetBase = 0, gmOffset = 0, indexOffset = 0, tailAlignCount = CeilAlignA2B(tailCount, perBlockCount);
    T x1, y1, z1;
    for (int64_t i = 0; i < bSize; i++) {
        gmOffsetBase = i * mSize * 3;
        for (int64_t j = 0, index = i * nSize * 3; j < nSize; j++, index += 3) {
            x1 = x1GM.GetValue(index);
            y1 = x1GM.GetValue(index + 1);
            z1 = x1GM.GetValue(index + 2);
            indexOffset = 0;
            gmOffset = 0;
            best1 = FLOAT_MAX;
            best2 = FLOAT_MAX;
            best3 = FLOAT_MAX;
            for (int64_t k = 0; k < loopTime; k++) {
                CopyInX2(gmOffsetBase + gmOffset, ubMaxProcCount);
                CopySplit(ubMaxGroupCount);
                ComputeNormal(x1, y1, z1, ubMaxGroupCount, indexOffset);
                gmOffset += ubMaxProcCount;
                indexOffset += ubMaxGroupCount;
            }
            if (tailCount) {
                CopyInX2(gmOffsetBase + gmOffset, tailAlignCount);
                CopySplit(tailGroupCount);
                ComputeNormal(x1, y1, z1, tailGroupCount, indexOffset);
            }
            distGM.SetValue(index, best1);
            distGM.SetValue(index + 1, best2);
            distGM.SetValue(index + 2, best3);
            idxGM.SetValue(index, bestIndex1);
            idxGM.SetValue(index + 1, bestIndex2);
            idxGM.SetValue(index + 2, bestIndex3);
        }
    }
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::CopyInX2(int64_t gmOffset, int64_t dataCount) {
    LocalTensor<T> xInLT = x2Queue.template AllocTensor<T>();
    DataCopy(xInLT, x2GM[gmOffset], dataCount);
    x2Queue.EnQue(xInLT);
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::CopySplit(int64_t groupCount) {
    LocalTensor<T> xLT = x2Queue.template DeQue<T>();
    int64_t tempCoef = CeilAlignA2B(groupCount, perBlockCount);
    LocalTensor<T> x1 = tempValLT;
    LocalTensor<T> y1 = tempValLT[tempCoef];
    LocalTensor<T> z1 = tempValLT[tempCoef * 2];
    pipe_barrier(PIPE_ALL);
    for (int64_t i = 0, index = 0; i < groupCount; i++) {
        x1.SetValue(i, xLT.GetValue(index++));
        y1.SetValue(i, xLT.GetValue(index++));
        z1.SetValue(i, xLT.GetValue(index++));
    }
    pipe_barrier(PIPE_ALL);
    x2Queue.FreeTensor(xLT);
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::ComputePerf(T x1, T y1, T z1, int64_t groupCount) {
    int64_t tempCoef = CeilAlignA2B(groupCount, perBlockCount);
    LocalTensor<T> x2 = tempValLT[tempCoef * 3];
    LocalTensor<T> y2 = tempValLT[tempCoef * 4];
    LocalTensor<T> z2 = tempValLT[tempCoef * 5];

    pipe_barrier(PIPE_ALL);
    Adds(x2, tempValLT[0], -x1, groupCount);
    Adds(y2, tempValLT[tempCoef * 1], -y1, groupCount);
    Adds(z2, tempValLT[tempCoef * 2], -z1, groupCount);

    Mul(x2, x2, x2, groupCount);
    Mul(y2, y2, y2, groupCount);
    Mul(z2, z2, z2, groupCount);

    Add(x2, x2, y2, groupCount);
    Add(x2, x2, z2, groupCount);

    ReduceMin(y2, x2, z2, groupCount, true);
    pipe_barrier(PIPE_ALL);
    best1 = y2.GetValue(0);
    bestIndex1 = *(reinterpret_cast<__ubuf__ int32_t*>(y2.GetPhyAddr(1)));
    x2.SetValue(bestIndex1, FLOAT_MAX);
    pipe_barrier(PIPE_ALL);
    ReduceMin(y2, x2, z2, groupCount, true);
    pipe_barrier(PIPE_ALL);
    best2 = y2.GetValue(0);
    bestIndex2 = *(reinterpret_cast<__ubuf__ int32_t*>(y2.GetPhyAddr(1)));
    x2.SetValue(bestIndex2, FLOAT_MAX);
    pipe_barrier(PIPE_ALL);
    ReduceMin(y2, x2, z2, groupCount, true);
    pipe_barrier(PIPE_ALL);
    best3 = y2.GetValue(0);
    bestIndex3 = *(reinterpret_cast<__ubuf__ int32_t*>(y2.GetPhyAddr(1)));
    pipe_barrier(PIPE_ALL);
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::CmpBest(T tempData, int32_t tempIndex) {
    if (tempData < best1) {
        best3 = best2;
        bestIndex3 = bestIndex2;
        best2 = best1;
        bestIndex2 = bestIndex1;
        best1 = tempData;
        bestIndex1 = tempIndex;
    } else if (tempData < best2) {
        best3 = best2;
        bestIndex3 = bestIndex2;
        best2 = tempData;
        bestIndex2 = tempIndex;
    } else if (tempData < best3) {
        best3 = tempData;
        bestIndex3 = tempIndex;
    }
}

template <typename T, int32_t BUFFER_NUM>
__aicore__ inline void ThreeNNND<T, BUFFER_NUM>::ComputeNormal(T x1, T y1, T z1, int64_t groupCount,
                                                               int64_t indexOffset) {
    int64_t tempCoef = CeilAlignA2B(groupCount, perBlockCount);
    LocalTensor<T> x2 = tempValLT[0];
    LocalTensor<T> y2 = tempValLT[tempCoef];
    LocalTensor<T> z2 = tempValLT[tempCoef * 2];

    pipe_barrier(PIPE_ALL);
    Adds(x2, x2, -x1, groupCount);
    Adds(y2, y2, -y1, groupCount);
    Adds(z2, z2, -z1, groupCount);

    Mul(x2, x2, x2, groupCount);
    Mul(y2, y2, y2, groupCount);
    Mul(z2, z2, z2, groupCount);

    Add(x2, x2, y2, groupCount);
    Add(x2, x2, z2, groupCount);

    T tempData[3];
    int32_t tempIndex[3];
    ReduceMin(y2, x2, z2, groupCount, true);
    pipe_barrier(PIPE_ALL);
    tempData[0] = y2.GetValue(0);
    tempIndex[0] = *(reinterpret_cast<__ubuf__ int32_t*>(y2.GetPhyAddr(1)));
    x2.SetValue(tempIndex[0], FLOAT_MAX);
    pipe_barrier(PIPE_ALL);
    ReduceMin(y2, x2, z2, groupCount, true);
    pipe_barrier(PIPE_ALL);
    tempData[1] = y2.GetValue(0);
    tempIndex[1] = *(reinterpret_cast<__ubuf__ int32_t*>(y2.GetPhyAddr(1)));
    x2.SetValue(tempIndex[1], FLOAT_MAX);
    pipe_barrier(PIPE_ALL);
    ReduceMin(y2, x2, z2, groupCount, true);
    pipe_barrier(PIPE_ALL);
    tempData[2] = y2.GetValue(0);
    tempIndex[2] = *(reinterpret_cast<__ubuf__ int32_t*>(y2.GetPhyAddr(1)));
    CmpBest(tempData[0], indexOffset + tempIndex[0]);
    CmpBest(tempData[1], indexOffset + tempIndex[1]);
    CmpBest(tempData[2], indexOffset + tempIndex[2]);
    pipe_barrier(PIPE_ALL);
}

}  // namespace ThreeNN

#endif  // THREE_N_N_N_D_H